# predict.py
import numpy as np
import torch
import matplotlib.pyplot as plt

# 导入自定义模块
from models.model import 丰度预测模型

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def predict_new_s(new_s):
    """
    预测新的物种存在向量s对应的丰度p
    new_s: 形状为(1, 100)的0-1向量（单个样本）
    """
    # 1. 加载模型
    model = 丰度预测模型().to(DEVICE)
    model.load_state_dict(torch.load("models/丰度预测模型.pth", map_location=DEVICE))
    model.eval()  # 推理模式

    # 2. 预处理输入（转换为张量）
    new_s_tensor = torch.tensor(new_s, dtype=torch.float32).to(DEVICE)

    # 3. 预测
    with torch.no_grad():
        predicted_p = model(new_s_tensor).cpu().numpy()  # 转回numpy数组

    return predicted_p


if __name__ == "__main__":
    # 示例：生成一个新的s向量（随机选50个物种存在）
    N = 100
    new_s = np.zeros((1, N))
    existing_species = np.random.choice(N, size=50, replace=False)  # 随机选50个物种
    new_s[0, existing_species] = 1  # 标记为存在

    # 预测对应的p
    predicted_p = predict_new_s(new_s)

    # 打印结果
    print("新的物种存在向量s中存在的物种索引：", existing_species)
    print("预测的丰度向量p（前10个物种）：", predicted_p[0, :10])

    # 可视化预测的丰度分布
    plt.figure(figsize=(10, 4))
    plt.bar(range(N), predicted_p[0], alpha=0.7)
    plt.xlabel("物种索引")
    plt.ylabel("预测丰度")
    plt.title("新群落的预测物种丰度分布")
    plt.savefig("predicted_p.png")
    plt.show()